import json
import glob
import os

import torch
from torch.utils.data import Dataset
from PIL import Image

from utils.preprocessText import preprocessText


class PS_Dataset(Dataset):
    def __init__(self, json_file, data_path, transformer, config, max_words=30):
        self.datas = json.load(open(json_file, 'r'))

        self.preprocess_text = preprocessText(config)
        self.data_path = data_path
        self.max_words = max_words
        self.transformer = transformer

    def __len__(self):
        return len(self.datas)

    def __getitem__(self, index):
        d = self.datas[index]

        # 预处理图像
        image_path = glob.glob(os.path.join(self.data_path, '%s_*.jpg' % d['id']))[0]
        image = Image.open(image_path).convert('RGB')
        image = self.transformer(image)

        # 预处理文本
        sentence = self.preprocess_text.pre_caption(d['text'], self.max_words)
        sentence = self.preprocess_text.replace_argot(sentence)

        # 处理标签
        label = [d['label'] - 1]
        label = torch.Tensor(label).long()
        return label, sentence, image

def statistic_class(dataset):
    statistic = {
        1: 0,
        2: 0,
        3: 0,
        4: 0
    }

    for i in dataset.indices:
        statistic[dataset.dataset.datas[i]['label']] += 1

    return statistic
